/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 */
#include "Segment.h"
#include "Common.h"
#include <atomic>
#include <cstring>
#include <iostream>

#include "Logger.h"

void Segment::init(int32_t poolId, int64_t segment_id, uint64_t offset,
                   uint8_t *memAddr, size_t memLength, uint64_t accessTime,
                   bool isShare, bool isFilled, std::string shmFilePath,
                   bool isSaved, int32_t writeRecordSize, int32_t startLayer,
                   int32_t endLayer, int32_t layerNum, int32_t layerCacheSize,
                   int64_t parentHashId, uint64_t localHashId) {
    _poolId         = poolId;
    _segmentId      = segment_id;
    _offset         = offset;
    _memAddr        = memAddr;
    _memLength      = memLength;
    _accessTime     = accessTime;
    _isShare        = isShare;
    _isFilled       = isFilled;
    _shmFilePath    = shmFilePath;

    _writeInRecord.store(writeRecordSize);

    _startLayer     = startLayer;
    _endLayer       = endLayer;
    _layerNum       = layerNum;
    _layerCacheSize = layerCacheSize;
    _isSaved        = isSaved;

    _hashId         = segment_id;
    _parentHashId   = parentHashId;
    _localHashId    = localHashId;
}


int32_t Segment::getPoolId() const {
  return _poolId; 
}


int64_t Segment::getSegmentId() const { 
  return _segmentId; 
}


uint64_t Segment::getOffset() const { 
  return _offset; 
}


uint8_t *Segment::getMemAddr() const { 
  return _memAddr; 
}


int32_t Segment::getSegmentSize() const { 
  return _memLength; 
}


uint64_t Segment::getAccessTime() const { 
  return _accessTime; 
}


bool Segment::getIsShare() const { 
  return _isShare; 
}


bool Segment::getIsFilled() const { 
  return _isFilled; 
}


std::string Segment::getShmFilePath() const { 
  return _shmFilePath; 
}


int32_t Segment::getWriteInRecord() const { 
  return _writeInRecord.load(); 
}


int32_t Segment::getStartLayer() const { 
  return _startLayer; 
}


int32_t Segment::getEndLayer() const { 
  return _endLayer; 
}


int32_t Segment::getLayerNum() const { 
  return _layerNum; 
}


int32_t Segment::getLayerCacheSize() const { 
  return _layerCacheSize; 
}


int64_t Segment::getHashId() const { 
  return _hashId; 
}


int64_t Segment::getParentHashId() const { 
  return _parentHashId; 
}


uint64_t Segment::getLocalHashId() const { 
  return _localHashId; 
}


const char *Segment::getVerifyHashId() const { 
  return _verifyHashId; 
}


bool Segment::getIsSaved() const { 
  return _isSaved; 
}


void Segment::setSegmentId(int64_t segmentId) { 
  _segmentId = segmentId; 
}


void Segment::setIsFilled(bool val) { 
  _isFilled = val; 
}


void Segment::setIsSaved(bool val) { 
  _isSaved = val; 
}


void Segment::setAccessTime(uint64_t timestamp) { 
  _accessTime = timestamp; 
}


void Segment::setWriteInRecord(int32_t val) { 
  _writeInRecord.store(val); 
}


void Segment::decWriteInRecord() { 
  _writeInRecord.fetch_sub(1); 
}


void Segment::setCacheInfo(CacheInfo &cacheInfo) {
    _startLayer     = cacheInfo.startLayer;
    _endLayer       = cacheInfo.endLayer;
    _layerNum       = cacheInfo.layerNum;
    _layerCacheSize = cacheInfo.layerCacheSize;
}


void Segment::setBlockHashId(BlockHashId &blockHashId) {
    _hashId       = blockHashId.hashId;
    
    _parentHashId = blockHashId.parentHashId;
    _localHashId  = blockHashId.localHashId;

    std::copy(blockHashId.verifyHashId,
            blockHashId.verifyHashId + VERIFY_HASH_LENGTH, 
            _verifyHashId
        );
}


#ifdef USE_GPU
int32_t gpuWriteIn(uint8_t *dstAddr, uint8_t *srcAddr, size_t size,
                   void *privData) {
  int32_t ret = KVC_OK;
  if (privData == nullptr) {
    ret = cudaMemcpy(dstAddr, srcAddr, size, cudaMemcpyDeviceToHost);
  } else {
    cudaStream_t stream = *(cudaStream_t *)privData;
    ret =
        cudaMemcpyAsync(dstAddr, srcAddr, size, cudaMemcpyDeviceToHost, stream);
  }
  if (ret != cudaSuccess) {
    LOG_ERROR_FMT("write from cuda failed, ret: {}", ret);
    ret = KVC_ERR_MEMFUNC_FAIL;
  }
  return ret;
}
#endif

#ifdef USE_NPU
int32_t npuWriteIn(uint8_t *dstAddr, size_t dstSize, uint8_t *srcAddr,
                   size_t srcSize, void *privData) {
  int32_t ret = KVC_OK;
  if (privData == nullptr) {
    LOG_ERROR_FMT("privData is nullptr");
    ret = aclrtMemcpy(dstAddr, dstSize, srcAddr, srcSize,
                      ACL_MEMCPY_DEVICE_TO_HOST);
  } else {
    aclrtStream stream = *(aclrtStream *)privData;
    ret = aclrtMemcpyAsync(dstAddr, dstSize, srcAddr, srcSize,
                           ACL_MEMCPY_DEVICE_TO_HOST, stream);
  }
  if (ret != ACL_SUCCESS) {
    LOG_ERROR_FMT("write from npu failed, ret: {}", ret);
    ret = KVC_ERR_MEMFUNC_FAIL;
  }
  return ret;
}
#endif

int32_t Segment::writeIn(DeviceType devType, void *buffer, uint64_t segOffset,
                         uint32_t size, void *privData) {
  if (_memLength < segOffset + size) {
    LOG_ERROR_FMT(" writeIn failed, segment mem size: {}, segOffset: {}, "
                  "buffer size: {}, shmOffset: {}",
                  _memLength, segOffset, size, _offset);
    return KVC_ERR_INVALID_ARG;
  }

  if (_isFilled) {
    LOG_ERROR_FMT("segment read only! memAddr: {}, size: {}, shmOffset: {}",
                  _memAddr, size, _offset);
    return KVC_ERR_MEM_REDAONLY;
  }

  int32_t ret = KVC_OK;
  switch (devType) {
  case DEVICE_TYPE_CPU:
    std::copy((uint8_t *)buffer, (uint8_t *)buffer + size,
              _memAddr + segOffset);
    break;
  case DEVICE_TYPE_GPU:
#if defined(USE_GPU)
    ret = gpuWriteIn((uint8_t *)_memAddr + segOffset, (uint8_t *)buffer, size,
                     privData);
    break;
#else
    LOG_ERROR_FMT("write from cuda failed, ret: {}", ret);
    throw std::runtime_error(
        "If use gpu, please add -DDEVICE=GPU to cmake args");
#endif
  case DEVICE_TYPE_NPU:
#if defined(USE_NPU)
    ret = npuWriteIn((uint8_t *)_memAddr + segOffset, _memLength - segOffset,
                     (uint8_t *)buffer, size, privData);
    break;
#else
    LOG_ERROR_FMT("If use npu, please add -DACL=1 to cmake args");
    throw std::runtime_error(
        "If use npu, please add -DDEVICE=NPU to cmake args");
#endif
  default:
    LOG_ERROR_FMT("Invalid device type");
    ret = KVC_ERR_NOT_SUPPORT;
    break;
  }

  return ret;
}

int32_t Segment::copyTo(DeviceType devType, void *buffer, uint64_t segOffset,
                        uint32_t size, void *privData) {
  if (_memLength < segOffset + size) {
    LOG_ERROR_FMT(
        "copyTo failed, segment mem size: {}, segOffset: {}, size: {}",
        _memLength, segOffset, size);
    return KVC_ERR_INVALID_ARG;
  }
  int32_t ret = KVC_OK;
  switch (devType) {
  case DEVICE_TYPE_CPU:
    std::copy(_memAddr + segOffset, _memAddr + segOffset + size,
              (uint8_t *)buffer);
    break;
  case DEVICE_TYPE_GPU:
#if defined(USE_GPU)
    if (privData == nullptr) {
      LOG_WARN_FMT("privData is nullptr");
      ret = cudaMemcpy(buffer, (char *)_memAddr + segOffset, size,
                       cudaMemcpyHostToDevice);
    } else {
      cudaStream_t stream = *(cudaStream_t *)privData;
      ret = cudaMemcpyAsync(buffer, (char *)_memAddr + segOffset, size,
                            cudaMemcpyHostToDevice, stream);
    }
    if (ret != KVC_OK) {
      LOG_ERROR_FMT("copy to cuda failed, ret: {}", ret);
      ret = KVC_ERR_MEMFUNC_FAIL;
    }
    break;
#else
    throw std::runtime_error(
        "If use gpu, please add -DDEVICE=GPU to cmake args");
#endif
  case DEVICE_TYPE_NPU:
#if defined(USE_NPU)
    if (privData == nullptr) {
      ret = aclrtMemcpy(buffer, _memLength - segOffset,
                        (char *)_memAddr + segOffset, size,
                        ACL_MEMCPY_HOST_TO_DEVICE);
    } else {
      aclrtStream stream = *(aclrtStream *)privData;
      ret = aclrtMemcpyAsync(buffer, _memLength - segOffset,
                             (char *)_memAddr + segOffset, size,
                             ACL_MEMCPY_HOST_TO_DEVICE, stream);
    }
    if (ret != ACL_SUCCESS) {
      LOG_ERROR_FMT("copy to npu failed, ret: {}", ret);
      ret = KVC_ERR_MEMFUNC_FAIL;
    }
    break;
#else
    throw std::runtime_error(
        "If use npu, please add -DDEVICE=NPU to cmake args");
#endif
  default:
    LOG_ERROR_FMT("Invalid device type: {}", devType);
    ret = KVC_ERR_NOT_SUPPORT;
    break;
  }
  return ret;
}